#! /usr/bin/env python
# -*- coding: utf-8 -*-

# *************************************************************
#  Copyright (c) Huoty - All rights reserved
#
#      Author: Huoty <sudohuoty@gmail.com>
#  CreateTime: 2016-10-08 10:46:31
# *************************************************************

"""加解密"""

import hashlib
from base64 import b64encode, b64decode

from Crypto import Random
from Crypto.Cipher import AES


def md5(data):
    m = hashlib.md5()
    m.update(data.encode("utf-8"))
    return m.hexdigest()

def sha1(data):
    return hashlib.sha1(data.encode("utf-8")).hexdigest()


class AESCipher(object):
    """AES 加密解密

    密钥 key 必须为 16（AES-128）, 24（AES-192） 或者 32（AES-256）Bytes 长度
    """
    def __init__(self, key):
        self.key = key if isinstance(key, bytes) else key.encode("utf-8")
        self.bs = AES.block_size

    def encrypt(self, text):
        # 若 text 为 str 类型, 则默认以 UTF-8 编码转化为 bytes 类型
        text = text if isinstance(text, bytes) else text.encode("utf-8")

        # Padding，由于 AES 是以 4×4 字节矩阵作为单位进行处理
        # 因此待加密数据必须是 16 的倍数，若不足 16 的倍数，则进行填充操作
        lack = self.bs - len(text) % self.bs  # 如果刚好为 16 的倍数则多加 16 位
        text = text + (lack * chr(lack)).encode("utf-8")

        # Initialization vector, 初始化向量
        # 用于避免相同的数据多次加密都产生相同的密文，最大长度为 16 字节
        iv = Random.new().read(self.bs)

        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        encrypted_text = cipher.encrypt(text)

        # 因为 AES 加密时候得到的字符串不一定是 ascii 字符集的，保存时候可能存在问题
        # 所以统一把加密后的字符串用 base64 转化
        return b64encode(iv + encrypted_text).decode('utf-8')

    def decrypt(self, text):
        text = b64decode(text)
        iv = text[:self.bs]
        cipher = AES.new(self.key, AES.MODE_CBC, iv)
        decrypted_text = cipher.decrypt(text[self.bs:])
        needless = ord(decrypted_text[len(decrypted_text)-1:])  # 注：decrypted_text[-1] 为 int 类型
        return decrypted_text[:-needless].decode('utf-8')


__all__ = ["md5", "sha1", "AESCipher"]


# Script starts from here

if __name__ == "__main__":
    pass
